package com.jack.leetcode.eachproblem_2020;

import java.util.*;

class RandomizedCollection {

    HashMap<Integer, Set<Integer>> map;
    List<Integer> list;
    int n;

    /** Initialize your data structure here. */
    public RandomizedCollection() {
        map = new HashMap<>();
        list = new ArrayList<>();
    }

    /** Inserts a value to the collection. Returns true if the collection did not already contain the specified element. */
    public boolean insert(int val) {
        Set<Integer> set = map.getOrDefault(val, new HashSet<>());
        set.add(n);
        map.put(val, set);
        list.add(val);
        n++;
        return set.size() == 1;
    }

    /** Removes a value from the collection. Returns true if the collection contained the specified element. */
    public boolean remove(int val) {
        if(map.containsKey(val)){
            Set<Integer> set = map.get(val);
            Set<Integer> last = map.get(list.get(n - 1));
            int cur = set.iterator().next();
            swap(list, cur, n - 1);
            list.remove(n - 1);
            set.remove(cur);
            if(set.isEmpty()){
                map.remove(val);
            }
            last.remove(n - 1);
            last.add(cur);
            n--;
            return true;
        }else{
            return false;
        }
    }

    private void swap(List<Integer> list, int cur, int last){
        Integer tmp = list.get(cur);
        list.set(cur, list.get(last));
        list.set(last, tmp);
    }

    /** Get a random element from the collection. */
    public int getRandom() {
        return list.get(new Random().nextInt(n));
    }
}

